import torch,os,pdb
from diffusers import FluxPriorReduxPipeline, FluxPipeline
from diffusers.utils import load_image

FLUX_REDUX='/home/shengjie/ckp/FLUX.1-Redux-dev'
FLUX='/data/models/FLUX___1-dev'

pipe_prior_redux = FluxPriorReduxPipeline.from_pretrained(
                                    FLUX_REDUX, 
                                    torch_dtype=torch.bfloat16).to("cuda")
pipe = FluxPipeline.from_pretrained(
    FLUX , 
    text_encoder=None,
    text_encoder_2=None,
    torch_dtype=torch.bfloat16
).to("cuda")

examples_dir = '/data/shengjie/style1/'
save_dir = '/data/shengjie/synthesis1/'

imagefiles = os.listdir(examples_dir)

# test_img = os.path.join(examples_dir,imagefiles[0])
# for test_steps in [10,15,20,30,40,50]:
#     # test_steps = 10
#     save_test_img = os.path.join(save_dir,
#                                 os.path.splitext(imagefiles[0])[0]+f'_step={test_steps}'+\
#                                     os.path.splitext(imagefiles[0])[1]
#                                 )
#     # pdb.set_trace()

#     image = load_image( test_img )
#     pipe_prior_output = pipe_prior_redux(image)


#     images = pipe(
#         guidance_scale=2.5,
#         num_inference_steps=test_steps,
#         generator=torch.Generator("cpu").manual_seed(0),
#         **pipe_prior_output,
#     ).images
#     images[0].save(save_test_img)